# 1. 导入依赖
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from utils.clusterAlgorithm.KMeansUtils import KMeansUtils

# 2. 加载数据集
init_data = pd.read_csv('../static/data/iris.csv')
x_params = "petal_length"
y_params = "petal_width"
label_params = "class"
unique_label = init_data[label_params].unique()
data = init_data[[x_params, y_params]].values.reshape(-1, 2)  # m x 2
label = init_data[label_params].values.reshape(-1, 1)  # m x 1
# 3. 展示原始数据
plt.figure(figsize=(8, 6))
plt.subplot(1, 2, 1)
for i_label in unique_label:
    plt.scatter(data[label[:, 0]  == i_label][:, 0], data[label[:, 0]  == i_label][:, 1], label=i_label)
plt.legend()
plt.title('Original Data known Label')
plt.subplot(1, 2, 2)
plt.scatter(data[:, 0], data[:, 1])
plt.title('Original Data unknown Label')
plt.show()

# 4. K-Means 聚类
k = 3   # 聚类数目
max_iter = 2000  # 最大迭代次数
kmeans_utils = KMeansUtils(data, k)
(centroids, distances) = kmeans_utils.train(max_iter) # centroids: k x 2, distances: m x 1

# 5. 展示聚类结果
plt.figure(figsize=(8, 6))
plt.subplot(1, 2, 1)
for i_label in unique_label:
    plt.scatter(data[label[:, 0] == i_label][:, 0], data[label[:, 0]  == i_label][:, 1], label=i_label)
plt.legend()
plt.title('K-Means Clustering Result-before')
plt.subplot(1, 2, 2)
for index in range(centroids.shape[0]):
    plt.scatter(data[distances[:,0]==index][:,0], data[distances[:,0]==index][:,1], label='Cluster '+str(index))
plt.scatter(centroids[:, 0], centroids[:, 1], marker='*', c='r', s=100, label='Cluster Center')
plt.legend()
plt.ylim(0, 2.7)
plt.title('K-Means Clustering Result-after')
plt.show()




